{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/LC1332/Chat-Haruhi-Suzumiya/blob/main/characters/novel_collecting/%E6%96%B0%E5%B0%8F%E8%AF%B4%E6%8A%BD%E5%8F%96_without_kor.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "今天试图去掉KOR\n",
        "\n",
        "考虑使用ChatHaruhi中的BaseLLM的"
      ],
      "metadata": {
        "id": "lCl4ZBtOxjVx"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "afxphm5ZXUdY"
      },
      "source": [
        "# 环境"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kxt1oviOXWcf"
      },
      "source": [
        "本脚本只需要在CPU环境运行\n",
        "\n",
        "如果你不在colab运行，你需要一个openai不封禁的ip"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Cc9nVmSHXWHq",
        "outputId": "00f66e5d-9bc6-4443-e4ad-81dffb8a8800"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m17.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m76.5/76.5 kB\u001b[0m \u001b[31m7.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.7/1.7 MB\u001b[0m \u001b[31m75.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.4/49.4 kB\u001b[0m \u001b[31m5.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for google-search-results (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ],
      "source": [
        "!pip -q install langchain openai google-search-results tiktoken\n",
        "!pip -q install markdownify\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ysEvTErI8_xb"
      },
      "source": [
        "# 新小说提取"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O9DMlJxQwWCA"
      },
      "source": [
        "Implemented by 李鲁鲁"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LFZTFVhzXgfv"
      },
      "source": [
        "检查下列变量是否被正确复制\n",
        "\n",
        "- [ ] openai的key（不要把自己的key 上传到github！）\n",
        "- [ ] 指定下载小说 wget后面的网址\n",
        "- [ ] 指定保存位置 save_folder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "krApyCoH3oJq"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kmb1oQFs3xsp",
        "outputId": "b5cf8f10-819e-4fb3-dfb5-0220c1fe77c2"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2023-09-04 11:26:33--  https://raw.githubusercontent.com/LC1332/Prophet-Andrew-Ng/main/langchain/%E7%A5%9E%E9%9B%95%E4%BE%A0%E4%BE%A3.txt\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 2928208 (2.8M) [text/plain]\n",
            "Saving to: ‘神雕侠侣.txt’\n",
            "\n",
            "神雕侠侣.txt        100%[===================>]   2.79M  --.-KB/s    in 0.01s   \n",
            "\n",
            "2023-09-04 11:26:45 (218 MB/s) - ‘神雕侠侣.txt’ saved [2928208/2928208]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "!wget https://raw.githubusercontent.com/LC1332/Prophet-Andrew-Ng/main/langchain/%E7%A5%9E%E9%9B%95%E4%BE%A0%E4%BE%A3.txt\n",
        "\n",
        "input_name  = '/content/神雕侠侣.txt'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zlrE4zpzXmgd",
        "outputId": "389850aa-d2cd-4602-fd7c-796eba7d2d05"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n",
            "注意，文件夹 /content/drive/MyDrive/StoryGPT/shendiaoxialv_extract 已经存在\n"
          ]
        }
      ],
      "source": [
        "# mount google drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "import os\n",
        "\n",
        "save_folder = \"/content/drive/MyDrive/StoryGPT/shendiaoxialv_extract\"\n",
        "\n",
        "if not os.path.exists(save_folder):\n",
        "    os.makedirs(save_folder)\n",
        "else:\n",
        "    print('注意，文件夹',save_folder,'已经存在')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58Arx3x788Ko"
      },
      "source": [
        "# 下载和切片天龙八部"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "18SGBCpkmJcN"
      },
      "source": [
        "这里你要给小说切一下章节（最好）\n",
        "\n",
        "如果你非常非常懒\n",
        "\n",
        "可以使用\n",
        "\n",
        "```python\n",
        "chapters = [raw_text]\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "oOIeYDaT4s_4"
      },
      "outputs": [],
      "source": [
        "import re\n",
        "\n",
        "raw_text = open(input_name, encoding='utf-8').read()\n",
        "\n",
        "# chapter_pattern = r'第\\d+章[^\\n]+\\n'\n",
        "\n",
        "# chapters = raw_text.split(chapter_pattern)\n",
        "\n",
        "# print(len(chapters))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAgEOI23mTgv"
      },
      "source": [
        "我觉得这段代码现在比正则表达式好用\n",
        "\n",
        "当然这个代码是针对 第xx章 xxxx 这种格式的\n",
        "\n",
        "如果有的时候你是 ---第xx章--\n",
        "\n",
        "你要改一下那个startswith那个 或者把前面的特殊符号给strip掉"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nggnYZtpaizy",
        "outputId": "3d277833-6118-4734-d3f7-40db8f933392"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "40\n"
          ]
        }
      ],
      "source": [
        "chapters = []\n",
        "chapter_contents = []\n",
        "\n",
        "for line in raw_text.split('\\n'):\n",
        "    Flag = False\n",
        "    if line.strip().startswith('第') and len(line)<100:\n",
        "        # 遇到章节标题,将之前章节内容添加到结果列表\n",
        "\n",
        "        head = line.strip()\n",
        "        # print(head)\n",
        "        head = head[:min(10,len(head))]\n",
        "        if head.find('章',1)>0 or head.find('回',1)>0:\n",
        "            # print(head)\n",
        "            Flag = True\n",
        "\n",
        "    if Flag:\n",
        "        if chapter_contents:\n",
        "            chapters.append('\\n'.join(chapter_contents))\n",
        "            chapter_contents = []\n",
        "        # 记录当前章节标题\n",
        "        # chapters.append(line)\n",
        "    else:\n",
        "        # 累积章节内容\n",
        "        chapter_contents.append(line)\n",
        "\n",
        "# 添加最后一个章节内容\n",
        "if chapter_contents:\n",
        "    chapters.append('\\n'.join(chapter_contents))\n",
        "\n",
        "print(len(chapters))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_wBzIPjhmhAc"
      },
      "source": [
        "看一看提取出来的章节对不对"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C2--yWclatnp",
        "outputId": "2a446600-0760-4b72-f7cd-47d629f78597"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0 “越女采莲秋水畔，窄\n",
            "1 武三娘正没做理会处，\n",
            "2 郭靖在舟中潜运伸功，\n",
            "3 郭靖摆脱众道纠缠，提\n",
            "4 杨过摔在山坡，滚入树\n",
            "5 小龙女从怀里取出一个\n",
            "6 杨过随着小龙女穿越甬\n",
            "7 杨过轻轻推开窗门，闪\n",
            "8 陆无双正自惶急，听他\n",
            "9 耶律齐道：“三妹，你\n",
            "10 杨过只奔出两步，突然\n",
            "11 次日杨过在厅上用过早\n",
            "12 金轮法王双眼时开时合\n",
            "13 当下陆家庄上重开筵席\n",
            "14 石阵外达尔巴和众蒙古\n",
            "15 杨过与陆无双听得冯铁\n",
            "16 次晨杨过醒来，走出石\n",
            "17 樊一翁见杨过折柳枝作\n",
            "18 杨过身子与水面相触的\n",
            "19 杨过本欲置身于这场是\n",
            "20 杨过正想拔出匕首，忽\n",
            "21 郭靖与杨过眼见无幸，\n",
            "22 杨过睡到中夜，忽然听\n",
            "23 杨过见天竺僧淡碧色的\n",
            "24 周伯通抬头见杆顶无旗\n",
            "25 小龙女眼见全真教群道\n",
            "26 郭靖走进房去带上了门\n",
            "27 杨过仍以右手空袖搂在\n",
            "28 郭芙叫道：“是我的小\n",
            "29 这段宁静平安也无多时\n",
            "30 绝情谷占地甚广，群山\n",
            "31 当黄蓉、一灯、郭芙等\n",
            "32 大宋理宗皇帝开庆元年\n",
            "33 眼前之人，正是杨过。\n",
            "34 杨过请得周伯通来和瑛\n",
            "35 次日英雄大宴续开。郭\n",
            "36 那和尚肩头在地下一靠\n",
            "37 那日郭襄见金轮法王猛\n",
            "38 一行人取道南下，沿路\n",
            "39 次日清晨，郭靖等一行\n"
          ]
        }
      ],
      "source": [
        "for i,ch in enumerate(chapters):\n",
        "    strip_ch = ch.strip()\n",
        "    print(i,strip_ch[:10])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7lYbFWh2h44T",
        "outputId": "3141e329-0758-4424-ee4a-970c74eba5d5"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "40\n",
            "\n",
            "　　“越女采莲秋水畔，窄袖轻罗，暗露双金钏。\n",
            "　　照影摘花花似面，芳心只共丝争乱。\n",
            "　　鸡尺溪头风浪晚，雾重烟轻，不见来时伴。\n",
            "　　隐隐歌声归掉远，离愁引看江南岸。”\n",
            "　　一阵轻柔婉转的歌声，飘在烟\n"
          ]
        }
      ],
      "source": [
        "print(len(chapters))\n",
        "print(chapters[0][:100])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "2GeitXNE69hg"
      },
      "outputs": [],
      "source": [
        "import tiktoken\n",
        "\n",
        "enc = tiktoken.get_encoding(\"cl100k_base\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "cellView": "form",
        "id": "LiG5bOg76-_-"
      },
      "outputs": [],
      "source": [
        "#@title  定义divide函数，用来切分超长文本\n",
        "def divide_str(s, sep=['\\n', '.', '。']):\n",
        "    mid_len = len(s) // 2  # 中心点位置\n",
        "    best_sep_pos = len(s) + 1  # 最接近中心点的分隔符位置\n",
        "    best_sep = None  # 最接近中心点的分隔符\n",
        "    for curr_sep in sep:\n",
        "        sep_pos = s.rfind(curr_sep, 0, mid_len)  # 从中心点往左找分隔符\n",
        "        if sep_pos > 0 and abs(sep_pos - mid_len) < abs(best_sep_pos -\n",
        "                                                        mid_len):\n",
        "            best_sep_pos = sep_pos\n",
        "            best_sep = curr_sep\n",
        "    if not best_sep:  # 没有找到分隔符\n",
        "        return s, ''\n",
        "    return s[:best_sep_pos + 1], s[best_sep_pos + 1:]\n",
        "\n",
        "\n",
        "def strong_divide(s):\n",
        "    left, right = divide_str(s)\n",
        "\n",
        "    if right != '':\n",
        "        return left, right\n",
        "\n",
        "    whole_sep = ['\\n', '.', '，', '、', ';', ',', '；',\\\n",
        "                 '：', '！', '？', '(', ')', '”', '“', \\\n",
        "                 '’', '‘', '[', ']', '{', '}', '<', '>', \\\n",
        "                 '/', '''\\''', '|', '-', '=', '+', '*', '%', \\\n",
        "               '$', '''#''', '@', '&', '^', '_', '`', '~',\\\n",
        "                 '·', '…']\n",
        "    left, right = divide_str(s, sep=whole_sep)\n",
        "\n",
        "    if right != '':\n",
        "        return left, right\n",
        "\n",
        "    mid_len = len(s) // 2\n",
        "    return s[:mid_len], s[mid_len:]\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-cmOV5a8mlNW"
      },
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xbvld6mG7A-J",
        "outputId": "80b6d419-f04c-4e4b-a690-8118a245b1da"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "959\n"
          ]
        }
      ],
      "source": [
        "#@title 以1500 token为限，切分chunk，输出总chunk数量\n",
        "\n",
        "max_token_len = 1500\n",
        "chunk_text = []\n",
        "\n",
        "\n",
        "for chapter in chapters:\n",
        "\n",
        "    split_text = chapter.split('\\n')\n",
        "\n",
        "    curr_len = 0\n",
        "    curr_chunk = ''\n",
        "\n",
        "    tmp = []\n",
        "\n",
        "    for line in split_text:\n",
        "        line_len = len(enc.encode( line ))\n",
        "\n",
        "        if line_len <= max_token_len - 5:\n",
        "            tmp.append(line)\n",
        "        else:\n",
        "            print('divide line with length = ', line_len)\n",
        "            path = [line]\n",
        "            tmp_res = []\n",
        "\n",
        "            while path:\n",
        "                my_str = path.pop()\n",
        "                left, right = strong_divide(my_str)\n",
        "\n",
        "                len_left = len(enc.encode( left ))\n",
        "                len_right = len(enc.encode( right ))\n",
        "\n",
        "                if len_left > max_token_len - 15:\n",
        "                    path.append(left)\n",
        "                else:\n",
        "                    tmp_res.append(left)\n",
        "\n",
        "                if len_right > max_token_len - 15:\n",
        "                    path.append(right)\n",
        "                else:\n",
        "                    tmp_res.append(right)\n",
        "\n",
        "            for line in tmp_res:\n",
        "                tmp.append(line)\n",
        "\n",
        "    split_text = tmp\n",
        "\n",
        "    for line in split_text:\n",
        "        line_len = len(enc.encode( line ))\n",
        "\n",
        "        if line_len > max_token_len:\n",
        "            print('warning line_len = ', line_len)\n",
        "\n",
        "        if curr_len + line_len <= max_token_len:\n",
        "            curr_chunk += line\n",
        "            curr_chunk += '\\n'\n",
        "            curr_len += line_len\n",
        "            curr_len += 1\n",
        "        else:\n",
        "            chunk_text.append(curr_chunk)\n",
        "            curr_chunk = line\n",
        "            curr_len = line_len\n",
        "\n",
        "    if curr_chunk:\n",
        "        chunk_text.append(curr_chunk)\n",
        "\n",
        "    # break\n",
        "\n",
        "print(len(chunk_text))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 51,
      "metadata": {
        "id": "XzB5ZszL7rcr",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7fb4004f-885b-4377-94f0-6533f956f68e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "　　小舟与那怪客相距数丈，陆无双年纪虽小，手上劲力竟自不弱，这一掷也是甚准。程英叫了声：“表妹！”待要阻止，己然不及，只见那莲蓬径往怪客脸上飞去。那怪客头一仰，已咬住莲蓬，也不伸手去拿，舌头卷处，咬住莲蓬便大嚼起来。五个少女见他竟不剥出莲子，也不怕苦涩，就这么连瓣连衣的吞吃，互相望了几眼，忍不住格格而笑，一面划船近前，走上岸来。\n",
            "　　程英走到那人身边，拉一拉他衣襟，道：“老伯怕，这样不好吃的。”\n",
            "　　从袋里取出一个莲蓬，劈开莲房，剥出十几颗莲子，再将莲子外的青皮撕开、取出莲子中苦味的芯儿，然后递在怪客手里。那怪客嚼了几口，但觉滋味清香鲜美，与适才所吃的大不相同，咧嘴向程英一笑，点了点头。程英又剥了几枚莲子递给他。那怪客将莲子抛入口中，一阵乱嚼，仰天说道：“跟我来？”\n",
            "\n",
            "　　说着大踏步向西便走。\n",
            "　　陆无双一拉程英的手，道：“表姊，咱们跟他去。”三个女伴胆小，忙道：“快回家去罢，别走远了惹你娘骂。”陆无双扁扁嘴扮个鬼脸，见那怪客走得甚快，说道：“你不来算啦。”放脱表姊的手，向前追去。程英与表妹一同出来玩耍，不能撇下她自归，只得跟去。那三个女伴虽比她们大了好几岁，但个个怕羞胆怯，只叫了几声，便见那怪客与程陆二人先后走入了桑树丛后。\n",
            "\n",
            "　　那怪客走得甚快，见程陆二人脚步小跟随不上，先还停步等了几次，到后来不耐烦起来，突然转身，长臂伸处，一手一个，将两个女孩儿挟在腋下，飞步而行，二女只听耳边风声飒然，路上的石块青草不住在眼前移动。陆无双害怕起来，叫道：“放下我，放下我！”那怪客哪里理她。反而走得更快了。陆无双仰起头来，张口往他手掌缘上猛力咬去。那怪客手掌一碰。只把她牙齿撞得隐隐生痛，陆无双只得松开牙齿，一张嘴可不闲着，拚命的大叫大嚷。程英却是默不作声。\n",
            "\n",
            "　　那怪客又奔一阵，将二人放下地来。当地是个坟场。程英的小脸吓成惨白，陆无双却涨得满脸通红。程英道：“老伯伯，我们要回家了，不跟你玩啦！”\n",
            "　　那怪客两眼瞪视着她，一言不发。程英见他目光之中流露出一股哀愁凄惋、自怜自伤的神色，不自禁的起了同情之心，轻轻道：“要是没人陪你玩，明天你再到湖边来，我剥莲子给你吃。”那怪客叹道：“是啊，十年啦，十年来都没人陪我玩。”突然间目现凶光，恶狠狠的道：“何沅君呢？何沅君到哪里去了？”\n",
            "\n",
            "　　程英见他突然间声色俱厉，心里害怕，低声道：“我……我……我不知道。”那怪客抓住她手臂，将她身子摇了几摇，低沉着嗓子道：“何沅君呢？”\n",
            "　　程英给他吓得几欲哭了出来，泪水在眼眶中滚来滚去，却始终没有流下。那怪客咬牙切齿的道：“哭啊，哭啊！你干么不哭？哼，你在十年前就是这样。\n",
            "\n"
          ]
        }
      ],
      "source": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 不使用Kor的dialogue抽取"
      ],
      "metadata": {
        "id": "I4dLN6qe1Bh2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content\n",
        "!rm -rf /content/Haruhi-2-Dev\n",
        "!git clone https://github.com/LC1332/Haruhi-2-Dev\n",
        "%cd /content/Haruhi-2-Dev\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NfBLhGMW1GZ3",
        "outputId": "16c506d6-1b53-438e-d876-f05b5a7e3f88"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content\n",
            "Cloning into 'Haruhi-2-Dev'...\n",
            "remote: Enumerating objects: 731, done.\u001b[K\n",
            "remote: Counting objects: 100% (38/38), done.\u001b[K\n",
            "remote: Compressing objects: 100% (30/30), done.\u001b[K\n",
            "remote: Total 731 (delta 19), reused 18 (delta 8), pack-reused 693\u001b[K\n",
            "Receiving objects: 100% (731/731), 105.69 MiB | 43.27 MiB/s, done.\n",
            "Resolving deltas: 100% (349/349), done.\n",
            "/content/Haruhi-2-Dev\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/Haruhi-2-Dev"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zG9raEUj54D1",
        "outputId": "1135c14f-307d-47f5-8080-c3fe90342bf3"
      },
      "execution_count": 52,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/Haruhi-2-Dev\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "然后我们新建一个LangChainGPT的类"
      ],
      "metadata": {
        "id": "JxvfGnMa1t4I"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q chromadb transformers"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wR1IrKo52APE",
        "outputId": "12e1f43f-19e0-4106-f9dc-9ef07b904bdd"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m418.3/418.3 kB\u001b[0m \u001b[31m7.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.5/7.5 MB\u001b[0m \u001b[31m77.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.1/3.1 MB\u001b[0m \u001b[31m82.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.4/58.4 kB\u001b[0m \u001b[31m7.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m59.5/59.5 kB\u001b[0m \u001b[31m7.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.4/5.4 MB\u001b[0m \u001b[31m99.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.9/5.9 MB\u001b[0m \u001b[31m105.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.8/7.8 MB\u001b[0m \u001b[31m103.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m67.3/67.3 kB\u001b[0m \u001b[31m8.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m593.7/593.7 kB\u001b[0m \u001b[31m52.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m268.8/268.8 kB\u001b[0m \u001b[31m28.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m54.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m67.0/67.0 kB\u001b[0m \u001b[31m7.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m46.0/46.0 kB\u001b[0m \u001b[31m4.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m6.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m428.8/428.8 kB\u001b[0m \u001b[31m35.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.1/4.1 MB\u001b[0m \u001b[31m74.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m55.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m129.9/129.9 kB\u001b[0m \u001b[31m14.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m86.8/86.8 kB\u001b[0m \u001b[31m10.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for chroma-hnswlib (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for pypika (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q transformers"
      ],
      "metadata": {
        "id": "loqnzJF_2c9C"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!ls"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bMTLsiLP2tUv",
        "outputId": "15188240-b67b-4811-c4e0-057435d01dc8"
      },
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "ChatHaruhi  data  LICENSE  notebook  Readme.md\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from ChatHaruhi.LangChainGPT import LangChainGPT\n",
        "\n",
        "llm_model = LangChainGPT()"
      ],
      "metadata": {
        "id": "FDXKbRr11yF3"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "smart_system_prompt = \"\"\"You are ChatGPT, a large language model trained by OpenAI.\n",
        "Knowledge cutoff: 2021-09\n",
        "Current date: 2023-03-15\"\"\"\n",
        "\n",
        "# llm_model.system_message( smart_system_prompt )"
      ],
      "metadata": {
        "id": "oWhbUqMm3CjF"
      },
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "task_prompt = \"\"\"Your task is to extract dialog text from each novel paragraph.\n",
        "For each paragraph, if there is dialog information in the sentence (dialog text is generally between quotation marks),\n",
        "infer the speaker and spoken text based on context, and output in the format <speaker>|<dialog text>.\n",
        "If there is no dialog in sentence, do not output anything.\"\"\"\n",
        "\n",
        "# llm_model.user_message( task_prompt )"
      ],
      "metadata": {
        "id": "ubnqXP-t3IMk"
      },
      "execution_count": 68,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(chunk_text[12])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RjSEIHqF6NbO",
        "outputId": "bee041f6-2162-4f16-fbd7-f4a0b9d5a26c"
      },
      "execution_count": 69,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "　　他见那女孩约莫比自己小着两三岁，人矮腿短，自己一发足便可追上，哪知他刚展开轻功，那女孩脚步好快，片刻间已奔出数丈，竟把他远远抛在后面，她再奔几步，站定身子，回头叫道：“哼，你追得着我么？”武修文道：”自然追得着。”立即提气急追。\n",
            "　　那女孩回头又跑，忽然向前疾冲，躲在一株松树后面。武修文随后跟来，那女孩瞧他跑得近了。斗然间伸出左足，往他小腿上绊上，武修文全没料到，登时向前跌出。他忙使个“铁树桩”想定住身子，那女孩右足又出，向他臀部猛力踢去。武修文一交直摔下去，鼻子刚好撞在一块小尖石上，鼻血流出；衣上点点斑斑的尽是鲜血。\n",
            "\n",
            "　　那女孩见血，不禁慌了。登时没做理会处，只想拔足逃走。忽然身后有人喝道：“芙儿.你又在欺侮人了。是不是？”那女孩并不回头，辩道：“谁说的？他自己摔交，管我甚么事？你可别跟我爹乱说，”武修文按住鼻子，其实也不很疼，只是见到满手鲜血，心下惊慌。他听得女孩与人说话，转过身来，见是个撑着铁拐的跛足老者。那人两鬓如霜，形容枯槁，双眼翻白，是个瞎子。\n",
            "\n",
            "　　只听他冷笑道：“你别欺我瞧不见，我甚么都听得清清楚楚。你这小妞儿啊，现下已经这样坏，大了瞧你怎么得了？”那女孩过去挽住他的手臂，央求道：“大公公，你别跟我爹爹说，好不好？他摔出了鼻血，你给他治治啊！”\n",
            "\n",
            "　　那老者踏上一步，左手抓住武修文手臂，右手伸指在他鼻旁“闻香穴”按了几按。武修文鼻血本已渐止，这么几按，就全然不流了，只觉那老者五根手指有如铁钳，又长又硬，紧紧抓着自己手臂，心中害怕起来，微微一挣，竟是动也不动，当下手臂一缩一圈，使出母亲所授的小擒拿手功夫，手掌打个半圈，向外逆翻。那老者没料到这小小孩童竟有如此巧妙手法，被他一翻之下，竟尔脱手，“噫”的一声轻呼，随即又抓住了他手腕。武修文运劲欲再挣扎，却怎么也挣不脱了。\n",
            "\n",
            "　　那老者道：“小兄弟别怕，你姓甚么？”武修文道：“我姓武。”那老者道：“你说话不是本地口音，从哪里来的？你爹妈呢？”说着放松了他手腕。武修文想起一晚没见爹娘，不知他两人怎样了，听他问起，险些儿便要哭出来。那女孩刮脸羞他，唱道：“羞羞羞，小花狗，眼圈儿红，要流油！”\n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "for ch in chunk_text:\n",
        "    if ch.find('那人',1)>0:\n",
        "        print(ch)\n",
        "        break"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qQZ5QQiL57FR",
        "outputId": "720544f1-a15b-4872-8f6a-2f083e2976ca"
      },
      "execution_count": 70,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "　　“越女采莲秋水畔，窄袖轻罗，暗露双金钏。\n",
            "　　照影摘花花似面，芳心只共丝争乱。\n",
            "　　鸡尺溪头风浪晚，雾重烟轻，不见来时伴。\n",
            "　　隐隐歌声归掉远，离愁引看江南岸。”\n",
            "　　一阵轻柔婉转的歌声，飘在烟水蒙蒙的湖面上。歌声发自一艘小船之中，船里五个少女和歌嘻笑，荡舟采莲。她们唱的曲子是北宋大词人欧阳修所作的“蝶恋花”词，写的正是越女采莲的情景，虽只寥寥六十字，但季节、时辰、所在、景物以及越女的容貌、衣着、首饰、心情，无一不描绘得历历如见，下半阕更是写景中有叙事，叙事中夹抒情，自近而远，余意不尽。欧阳修在江南为官日久，吴山越水，柔情蜜意，尽皆融入长短句中。宋人不论达官贵人，或是里巷小民，无不以唱词为乐，是以柳永新词一出，有井水处皆歌，而江南春岸折柳，秋湖采莲，随伴的往往便是欧词。\n",
            "\n",
            "　　时当南宋理宗年间，地处嘉兴南湖。节近中秋，荷叶渐残，莲肉饱实。这一阵歌声传入湖边一个道姑耳中。她在一排柳树下悄立已久，晚风拂动她杏黄色道袍的下摆，拂动她颈中所插拂尘的万缕柔丝，心头思潮起伏，当真亦是“芳心只共丝争乱”。只听得歌声渐渐远去，唱的是欧阳修另一首“蝶恋花”词，一阵风吹来，隐隐送未两句：“风月无情人暗换，旧游如梦空肠断……”\n",
            "\n",
            "　　歌声甫歇，便是一阵格格娇笑。\n",
            "　　那道姑一声长叹，提起左手，瞧着染满了鲜血的手掌，喃喃自语：“那又有甚么好笑？小妮子只是瞎唱，浑不解同中相思之苦、惆怅之意。”\n",
            "　　在那道姑身后十余丈处，一个青袍长须的老者也是一直悄立不动，只有当“风月无情人暗换，旧游如梦空肠断”那两句传到之时，发出一声极轻极轻的叹息。\n",
            "　　小船在碧琉璃般的湖面上滑过，舟中五个少女中三人十五六岁上下，另外两个都只九岁，两个幼女是中表之亲，表姊姓程，单名一个英字，表妹姓陆，名无双，两人相差半岁。\n",
            "　　三个年长少女唱着歌儿，将小舟从荷叶丛中荡将出来。程英道：“表妹你瞧，这位老伯伯还在这儿，”说着伸手指向垂柳下的一人。\n",
            "　　那人满头乱发，胡须也是蓬蓬松松如刺猬一般，须发油光乌黑，照说年纪不大，可是满脸皱纹深陷，却似七八十岁老翁。身穿蓝布直缀，颈中挂着个婴儿所用的锦缎围涎，围涎上绣着幅花猫扑蝶图，已然陈旧破烂。\n",
            "\n",
            "　　陆无双道：“这怪人在这儿坐了老半天啦，怎么动也不动？”程英道：“别叫怪人，要叫‘老伯伯’。你叫他怪人，他要生气的。”陆无双笑道：“他还不怪吗？这么老了，头颈里却挂了个围涎。他生了气，要是胡子都翘了起来，那才好看呢。”从小舟中拿起一个莲蓬，往那人头上掷去。\n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "ex_ins = []\n",
        "ex_outs = []\n",
        "\n",
        "example_in = \"\"\"\\n###\\n柯镇恶大怒，叫道：“就算打你不死，累也累死了你。”绕钟来追。\n",
        "　　杨过见二人绕着铁钟兜圈子，时候一长，义父必定气力不加，眼见情势危急，忽然心生一计\n",
        "杨过,“看，飞机”\"\"\"\n",
        "\n",
        "example_out = \\\n",
        "\"\"\"\n",
        "role|dialogue\n",
        "柯镇恶|“就算打你不死，累也累死了你。”\n",
        "杨过|“看，飞机”\n",
        "\"\"\"\n",
        "\n",
        "ex_ins.append(example_in)\n",
        "ex_outs.append(example_out)"
      ],
      "metadata": {
        "id": "6d8NKXEd4LdO"
      },
      "execution_count": 71,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "example_in = \"\"\"\\n###\\n张十五见他身材魁梧，浓眉大眼，便道：“正是。”那大汉道：“小弟作东，请先生去饮上三杯如何”张十五大喜，说道：“素不相识，怎敢叨扰”\n",
        "\n",
        "    那大汉笑道：“喝上三怀，那便相识了。我姓郭，名叫郭啸天。”\"\"\"\n",
        "\n",
        "example_out = \\\n",
        "\"\"\"\n",
        "role|dialogue\n",
        "张十五|“正是。”\n",
        "郭啸天|“小弟作东，请先生去饮上三杯如何”\n",
        "张十五|“素不相识，怎敢叨扰”\n",
        "郭啸天|“喝上三怀，那便相识了。我姓郭，名叫郭啸天。”\n",
        "\"\"\"\n",
        "\n",
        "ex_ins.append(example_in)\n",
        "ex_outs.append(example_out)\n",
        "\n"
      ],
      "metadata": {
        "id": "u3JG1OiQ6u7i"
      },
      "execution_count": 72,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "import os\n",
        "\n",
        "key = \"sk-VvF4zQ\"\n",
        "key_bytes = key.encode()\n",
        "os.environ[\"OPENAI_API_KEY\"] = key_bytes.decode('utf-8')"
      ],
      "metadata": {
        "id": "XSN0e_Kv9S52"
      },
      "execution_count": 73,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import csv\n",
        "import json\n",
        "\n",
        "def csv_to_json(text):\n",
        "    lines = []\n",
        "    for line in text.splitlines():\n",
        "        if '|' in line:\n",
        "            lines.append(line)\n",
        "\n",
        "    reader = csv.reader(lines, delimiter='|')\n",
        "\n",
        "    result = []\n",
        "    for row in reader:\n",
        "        if row[0] == 'role' and row[1] == 'dialogue':\n",
        "            continue\n",
        "        result.append({\n",
        "            \"role\": row[0],\n",
        "            \"dialogue\": row[1]\n",
        "        })\n",
        "\n",
        "    return result\n",
        "\n",
        "def extract_dialogue(input_text):\n",
        "    llm_model = LangChainGPT()\n",
        "\n",
        "    # llm_model.system_message( smart_system_prompt.encode().decode('utf-8') )\n",
        "\n",
        "    # llm_model.user_message( task_prompt.encode().decode('utf-8')  )\n",
        "    llm_model.system_message( task_prompt.encode().decode('utf-8')  )\n",
        "\n",
        "    for inp,out in zip(ex_ins, ex_outs):\n",
        "        llm_model.user_message(inp.encode().decode('utf-8') )\n",
        "        llm_model.ai_message(out.encode().decode('utf-8') )\n",
        "\n",
        "    input_text = '\\n###\\n' + input_text\n",
        "    llm_model.user_message(input_text.encode().decode('utf-8') )\n",
        "\n",
        "    response = llm_model.get_response()\n",
        "\n",
        "    json_result = csv_to_json(response)\n",
        "\n",
        "    return json_result\n",
        "\n",
        "\n",
        "response = extract_dialogue( chunk_text[i+1])\n",
        "print(response)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NHp0vGtB7Jeu",
        "outputId": "484aca57-3ca7-4d32-955f-56f507ac8d2e"
      },
      "execution_count": 74,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[{'role': '柯镇恶', 'dialogue': '“就算打你不死，累也累死了你。”'}]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "response = extract_dialogue( chunk_text[666])\n",
        "print(response)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jn94wLWU99qb",
        "outputId": "9590910f-b9d1-4d92-c6aa-d4ebc6a2cce2"
      },
      "execution_count": 75,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[{'role': '杨过', 'dialogue': '“啊哟！险些误了大事！”'}, {'role': '小龙女', 'dialogue': '“怎么？”'}, {'role': '杨过', 'dialogue': '“咱们练到紧要关头，要是这小鬼头突然叫嚷起来，岂不糟糕！”'}, {'role': '小龙女', 'dialogue': '“好险！”'}, {'role': '杨过', 'dialogue': '“你全身三十六处大穴尽数冲开，我瞧快则十日，慢须半月。本来这么多的时日之中，免不了有外物分心，但这古墓与尘世隔绝，当真是天下最好不过之地，便是最幽静的荒山穷谷，也总会有清风明月、鸟语花香扰人心神。”'}, {'role': '小龙女', 'dialogue': '“我这伤是全真道人打的，但全真教的祖师爷造了墓室、备了寒玉床，供我安安静静的休息，回复安康，他们的功罪也足以相抵了。”'}, {'role': '杨过', 'dialogue': '“那金轮法王呢？咱们可饶他不得。”'}, {'role': '小龙女', 'dialogue': '“只要我能活着，你还有甚么不满足的么？”'}, {'role': '杨过', 'dialogue': '“你说得是。这次你伤好了，咱们永远不再跟人动手。老天爷待咱们这么好！唉。”'}, {'role': '小龙女', 'dialogue': '“咱们到南方去，种几亩田，养些小鸡小鸭……”'}, {'role': '杨过', 'dialogue': '（无对话）'}, {'role': '小龙女', 'dialogue': '（无对话）'}, {'role': '杨过', 'dialogue': '（无对话）'}]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "response = extract_dialogue( chunk_text[666])\n",
        "print(response)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1mBnZjHkAliL",
        "outputId": "4843b710-38bc-4522-ed67-9fd0c60c10a2"
      },
      "execution_count": 59,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[{'role': '柯镇恶', 'dialogue': '“秦王鞭石”'}, {'role': '柯镇恶', 'dialogue': '“柠伏药叉”'}, {'role': '杨过', 'dialogue': '（心惊肉跳）有心要上前相助义父，却自知武艺低微，只有送死的份儿。'}, {'role': '柯镇恶', 'dialogue': '（连续三杖击中欧阳锋）'}, {'role': '欧阳锋', 'dialogue': '（受击后）暗暗叫苦'}, {'role': '柯镇恶', 'dialogue': '（与欧阳锋滚倒地上）'}, {'role': '欧阳锋', 'dialogue': '（抓住柯镇恶胸口）'}, {'role': '柯镇恶', 'dialogue': '（反揪欧阳锋）'}, {'role': '柯镇恶', 'dialogue': '“砰”的一声将欧阳锋打了个筋斗'}, {'role': '欧阳锋', 'dialogue': '（迷糊中）将尖刀往敌人掷去'}, {'role': '柯镇恶', 'dialogue': '（闪身避开刀）'}]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OCKTCEB29DlH"
      },
      "source": [
        "# 编辑新的Kor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F0EawLOADbnj"
      },
      "source": [
        "# Bullet形式摘要抽取"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 76,
      "metadata": {
        "id": "ANBKqeO6CXFP"
      },
      "outputs": [],
      "source": [
        "system_prompt = \"\"\"\n",
        "Summarize the key points of the following text in a concise way, using bullet points.\n",
        "\"\"\"\n",
        "\n",
        "q_example = \"\"\"###\n",
        "Text:\n",
        "洪七公、周伯通、郭靖、黄蓉四人乘了小船，向西驶往陆地。郭靖坐在船尾扳桨，黄蓉不住向周伯通详问骑鲨游海之事，周伯通兴起，当场就要设法捕捉鲨鱼，与黄蓉大玩一场。\n",
        "郭靖见师父脸色不对，问道：“你老人家觉得怎样”洪七公不答，气喘连连，声息粗重。他被欧阳锋以“透骨打穴法”点中之后，穴道虽已解开，内伤却又加深了一层。黄蓉喂他服了几颗九花玉露丸，痛楚稍减，气喘仍是甚急。\n",
        "老顽童不顾别人死活，仍是嚷着要下海捉鱼，黄蓉却已知不妥，向他连使眼色，要他安安静静的，别吵得洪七公心烦。周伯通并不理会，只闹个不休。黄蓉皱眉道：“你要捉鲨鱼，又没饵引得鱼来，吵些甚么”\n",
        "\n",
        "Summarize in BULLET POINTS form:\n",
        "\"\"\"\n",
        "\n",
        "a_example = \"\"\"\n",
        "- 洪七公等四人乘船西行,洪七公因受内伤加重而气喘不止\n",
        "- 周伯通要捉鲨鱼玩,被黄蓉阻止以免掀翻小船\n",
        "\"\"\"\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 77,
      "metadata": {
        "id": "x9oE6AYhGTCu"
      },
      "outputs": [],
      "source": [
        "from langchain.chat_models import ChatOpenAI\n",
        "from langchain.prompts.chat import (\n",
        "    ChatPromptTemplate,\n",
        "    SystemMessagePromptTemplate,\n",
        "    AIMessagePromptTemplate,\n",
        "    HumanMessagePromptTemplate,\n",
        ")\n",
        "from langchain.schema import (\n",
        "    AIMessage,\n",
        "    HumanMessage,\n",
        "    SystemMessage\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 78,
      "metadata": {
        "id": "VS-1K7mRGTbR"
      },
      "outputs": [],
      "source": [
        "llm = ChatOpenAI(\n",
        "    model_name=\"gpt-3.5-turbo\",\n",
        "    temperature=0\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 79,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ar_nQWRoGXJL",
        "outputId": "b0d1fb5c-bcbe-433d-8034-bb7327989117"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "- 杨过被金杵逼退，躲进厅角\n",
            "- 杨过使用九阴真经中的武功跳过金杵，救了自己\n",
            "- 郭靖与金轮法王交手，杨过也使用九阴真经中的武功躲过金杵\n",
            "- 杨过的九阴真经引起黄蓉的好奇，她想知道郭靖是否教过他\n",
            "- 黄蓉发现丐帮叛徒彭长老在蒙古武士群中，提醒杨过使用移魂大法\n"
          ]
        }
      ],
      "source": [
        "messages = [SystemMessage( content = system_prompt),\n",
        "            HumanMessage( content = q_example),\n",
        "            AIMessage( content = a_example)]\n",
        "\n",
        "sel_id = 300\n",
        "\n",
        "new_input = f\"\"\"###\n",
        "Text:\n",
        "{chunk_text[sel_id]}\n",
        "\n",
        "Summarize in BULLET POINTS form:\"\"\"\n",
        "\n",
        "messages.append( HumanMessage(content = new_input) )\n",
        "\n",
        "response = llm( messages )\n",
        "\n",
        "print(response.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TfJ0Fnl_Qyhr"
      },
      "source": [
        "理论上输出是这样的\n",
        "\n",
        "\n",
        "- 段誉和朱丹臣讨论诗词,表达友谊\n",
        "- 木婉清不懂诗词,觉得两人说书呆子话\n",
        "- 朱丹臣礼貌对木婉清,想请两人回段家\n",
        "- 段誉担心鳄神,同意回家,让朱丹臣护送\n",
        "- 木婉清说朱丹臣话太多,自己是山野女子不懂\n",
        "- 朱丹臣表示会简化自己的文词,不让木婉清觉得难懂\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 80,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WfA6NoYNQxXZ",
        "outputId": "c72f2165-0fdb-4ec4-d351-5d8bf714cac3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/drive/MyDrive/StoryGPT/shendiaoxialv_extract\n"
          ]
        }
      ],
      "source": [
        "print(save_folder)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SM63hA11HQTa",
        "outputId": "71fe217a-6ff1-43a4-fc60-6536c75839e5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1NLzjtsZOHiO"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hCW42kX9O5pO",
        "outputId": "2f5ae884-68c4-4837-b990-18d3e0db860a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  1%|          | 9/959 [00:42<1:31:36,  5.79s/it]"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import json\n",
        "from tqdm import tqdm\n",
        "\n",
        "# save_folder = \"/content/drive/MyDrive/GPTData/weixiaobao_extract\"\n",
        "\n",
        "for i in tqdm(range(len(chunk_text))):\n",
        "    save_name = os.path.join(save_folder, f\"{i}_dialogue.txt\")\n",
        "\n",
        "    if not os.path.exists(save_name) or os.path.getsize(save_name) < 0:\n",
        "        if os.path.exists(save_name):\n",
        "            print('re-generate dialogue id = ', i)\n",
        "\n",
        "        dialogue_response = extract_dialogue( chunk_text[i] )\n",
        "\n",
        "        with open(save_name, 'w', encoding='utf-8') as f:\n",
        "            for chat in dialogue_response:\n",
        "                json_str = json.dumps(chat, ensure_ascii=False)\n",
        "                f.write(json_str+\"\\n\")\n",
        "\n",
        "    save_name_sum = os.path.join(save_folder, f\"{i}_sum.txt\")\n",
        "\n",
        "    if not os.path.exists(save_name_sum) or os.path.getsize(save_name_sum) < 5:\n",
        "        if os.path.exists(save_name_sum):\n",
        "            print('re-summarize id = ',i )\n",
        "        #dealing with summarize\n",
        "        messages = [SystemMessage( content = system_prompt),\n",
        "                HumanMessage( content = q_example),\n",
        "                AIMessage( content = a_example)]\n",
        "\n",
        "        new_input = f\"###\\nText:\\n{chunk_text[ i ]}\\nSummarize in BULLET POINTS form:\"\n",
        "\n",
        "        messages.append( HumanMessage(content = new_input) )\n",
        "\n",
        "        summarize_response = llm( messages ).content\n",
        "\n",
        "        with open(save_name_sum, 'w', encoding='utf-8') as f:\n",
        "            f.write( summarize_response )\n",
        "\n",
        "    raw_text_save_name = os.path.join(save_folder, f\"{i}_raw.txt\")\n",
        "\n",
        "    if not os.path.exists(raw_text_save_name) or os.path.getsize(raw_text_save_name) < 5:\n",
        "        with open(raw_text_save_name, 'w', encoding='utf-8') as f:\n",
        "            f.write( chunk_text[i] )\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gbSAvk4mivfw",
        "outputId": "1b4ff050-58f7-4bf1-9697-f0cca809cb52"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "    游坦之乘着众人扰攘之际，和全冠清低声商议了一阵，又朗声道：“我大宋国步艰危，江湖同道却又不能齐心合力，以至时受番帮欺压。因此丐帮主张立一位武林盟主，大伙儿听奉号令，有什么大事发生，便不致乱成一团了。玄慈方丈，你赞不赞成？”    玄慈缓缓地道：“庄帮主的话，倒也言之成理。但老衲有一事不解，却要请教。”游坦之道：“什么事？”玄慈道：“庄帮主已拜丁先生为师，算是星宿派门人了，是也不是？”游坦之道：“这个……这是我自己的事，与你无关。”玄慈道：“星宿派乃西域门派，非我大宋武林同道。我大宋立不立武林盟主，可与星宿派无涉。就算中原武林同道要推举一位盟主，以便统筹事功，阁下是星宿派门人，却也不便参与了。”\n",
            "    众英雄纷纷说道：“不错！”“少林方丈之言甚是。”“你是番邦门派的走狗奴才，怎可妄想做我中原武林的盟主？。”\n",
            "    游坦之无言可答，向丁春秋望望，又向全冠清瞧瞧，盼望他们出言解围。\n",
            "    丁春秋咳嗽一声，说道：“少林方丈言之差矣！老夫乃山东曲阜人氏，生于圣人之邦，星宿派乃老夫一手创建，怎能说是西域番邦的门派？星宿派虽居处西域，那只不过是夫夫暂时隐居之地。你说星宿派是番邦门派，那么孔夫子也是番邦人氏了，可笑啊可笑！说到西域番邦，少林武功源于天竺达摩祖师，连佛教也是西域番邦之物，我看少林派才是西域的门派呢！”此言一出，玄慈和群雄都感不易抗辩。\n",
            "    全冠清朗声道：“天下武功，源流难考。西域武功传于中土者有之，中土武功传于西域者亦有之。我帮庄帮主乃中土人氏，丐帮素为中原门派，他自然是中原武林的领袖人物。玄慈方丈，今日之事，当以武功强弱定胜负，不以言辞舌辩定输赢。丐帮与少林派到底谁强谁弱，只须你们两位首领出手较量，高下立判，否则便是说上半天，又有何益？倘若你有自知之明，不是敝帮庄帮主的敌手，只须甘拜下风，推戴我庄帮主为武林盟主，倒也不是非出手不可的。”这几句话，显然认定玄慈是明知不敌，胆怯推诿。\n",
            "    玄慈向前走了几步，说道：“庄帮主，你既非要老衲出手不可，老衲若再顾念贵帮和敝帮数百年的交情，坚不肯允，倒是对贵帮不敬了。”眼光向群雄缓缓掠过，朗声道：“天下英雄，今日人人亲眼目睹，我少林派决无与丐帮争雄斗胜之意，实是丐帮帮主步步见逼，老衲退无可退，避无可避。”\n",
            "    群雄纷纷说道：“不错，咱们都是见证，少林派并无丝毫理亏之处。”\n",
            "    游坦之只是挂念着阿紫的安危，一心要尽快杀了玄慈，好向丁春秋交差，大声说道：“比武较量，强存弱亡，说不上谁理亏不理亏，快快上来动手吧！”\n",
            "\n"
          ]
        }
      ],
      "source": [
        "print(chunk_text[984])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XmB_H2Z4CGLr"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}